# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: jhkim

import datetime
import fire
import os
import pandas
import re
import json
import time
from collections import defaultdict

from coin.support.market_quotes.logic.market_status_util import \
  query_history_kline_list
import cc.appcoin2.strategy.lm.kline_generator as klg

pit_pi_path_base = "/remote/iosg/home-2/public/bot-coin/coin_product_info/data"

def get_pi_paths_by_exchange_date(market_type, exchange, date, days):
  pi_name_map = {
      ("Futures", "Binance"): ["Futures.Binance.v1.json", "Futures.Binance.v1-delivery.json"],
      ("Futures", "Huobi"): ["Futures.Huobi.v1-linear-swap.json"],
      ("Futures", "Okex"): ["Futures.Okex.v5-swap.json", "Futures.Okex.v5.json"],
      ("Futures", "Ftx"): ["Futures.Ftx.v1.json"],
      ("Futures", "Dydx"): ["Futures.Dydx.v3.json"],
      ("Futures", "Bybit"): ["Futures.Bybit.v2.json", "Futures.Bybit.v2-linear.json",
                             "Futures.Bybit.v3.json", "Futures.Bybit.v3-linear.json"],
      ("Futures", "Bitget"): ["Futures.Bitget.v1.json"],
      ("Futures", "Phemex"): ["Futures.Phemex.v1.json"],
      ("Futures", "Kucoin"): ["Futures.Kucoin.v1.json"],
      ("Futures", "Deribit"): ["Futures.Deribit.v2.json"],
      ("Spot", "Binance"): ["Spot.Binance.v1.json"],
      ("Spot", "Huobi"): ["Spot.Huobi.v1.json"],
      ("Spot", "Okex"): ["Spot.Okex.v5.json"],
      ("Spot", "Ftx"): ["Spot.Ftx.v1.json"],
      ("Spot", "Gdax"): ["Spot.Gdax.v1.json"],
      ("Spot", "Upbit"): ["Spot.Upbit.v1.json"],
      ("Spot", "Bithumb"): ["Spot.Bithumb.v2.json"],
      ("Spot", "Bitbank"): ["Spot.Bitbank.v1.json"],
      ("Spot", "Bybit"): ["Spot.Bybit.v1.json", "Spot.Bybit.v3.json"],
      ("Spot", "Btcturk"): ["Spot.Btcturk.v1.json"],
      ("Spot", 'Gateio'): ["Spot.Gateio.v4.json"],
      ("Spot", 'Mexc'): ["Spot.Mexc.v3.json"],
      ("Spot", 'Kucoin'): ["Spot.Kucoin.v1.json"],
      ("Spot", 'Kraken'): ["Spot.Kraken.v1.json"],
      ("Spot", 'Korbit'): ["Spot.Korbit.v1.json"],
      ("Spot", "Bitget"): ["Spot.Bitget.v1.json"],
      ("Spot", "Phemex"): ["Spot.Phemex.v1.json"],
      ("Spot", "Coingecko"): ["Spot.Coingecko.v3.json"],
  }
  pit_pi_paths = []
  date_range = []
  if date is not None:
    date_range = pandas.date_range(date - datetime.timedelta(days=days), date + datetime.timedelta(days=days)).sort_values(ascending=False)
  for d in date_range:
    d = d.strftime("%Y%m%d")
    path = os.path.join(pit_pi_path_base, d)
    if (market_type, exchange) not in pi_name_map:
      continue
    for pi_name in pi_name_map[(market_type, exchange)]:
      full_path = os.path.join(path, pi_name)
      if os.path.exists(full_path):
        pit_pi_paths.append(full_path)

  return pit_pi_paths

def get_symbols_from_adjacent_pis(market_type, exchange, append_contract_value=False, futures_suffix=".PERPETUAL", date=None, days=0):
  pi_paths = get_pi_paths_by_exchange_date(market_type, exchange, date, days)
  print(pi_paths)
  pi_infos = list()
  pi_symbols = set()
  for pi_path in pi_paths:
    with open(pi_path, "r", encoding="utf8") as fin:
      pi_info_each = json.loads(fin.read())["product_infos"]
      for pi_info in pi_info_each:
        if pi_info['symbol'] not in pi_symbols:
          pi_infos.append(pi_info)
          pi_symbols.add(pi_info['symbol'])
  symbols_with_pi = list()
  for _pi in pi_infos:
    symbols_with_pi.append({
      'symbol': _pi["symbol"],
      'contract_value': _pi.get("contract_value", 1.0),
      'is_inverse': _pi.get("is_inverse", False),
    })
  if market_type == "Futures":
    symbols_with_pi = list(filter(lambda u: u['symbol'].endswith(futures_suffix), symbols_with_pi))
  print(f"symbols from {pi_path}, len: {len(symbols_with_pi)}, head: {symbols_with_pi[:5]}")
  if append_contract_value:
    return symbols_with_pi
  return list(map(lambda u: u['symbol'], symbols_with_pi))

def generate_klineh5(trading_date: str, market: str, exchange: str, pattern='.*USD.*', coin_product_info_path=None, output_path=None):
  global pit_pi_path_base
  if coin_product_info_path is not None:
    pit_pi_path_base = os.path.join(coin_product_info_path, "data")
  interval_klineh5_root = "/remote/iosg/data-2/buckets/kline.derived.interval_h5"
  if output_path is not None:
    interval_klineh5_root = output_path
  trading_date = str(trading_date)
  trading_dates = []
  if "-" in trading_date:
    trading_dates = pandas.date_range(trading_date.split("-")[0], trading_date.split("-")[1])
  else:
    trading_dates = [datetime.datetime.strptime(trading_date, "%Y%m%d")]
  market = str(market)
  exchange = str(exchange)
  print(trading_date)

  freq_dict = defaultdict(lambda: "2H")
  freq_dict.update({
      "Binance": "12H",
      "Okex": "3H",
      "Bybit": "12H",
  })
  freq = freq_dict[exchange]

  for date_today in trading_dates:
    print(date_today.strftime("%Y%m%d"))
    dfdict = {}
    nodes = pandas.date_range(date_today, date_today + datetime.timedelta(days=1), freq=freq).to_pydatetime().tolist()
    ranges = list(zip(nodes[:-1], nodes[1:]))
    me = [market, exchange]
    symbols = get_symbols_from_adjacent_pis(me[0], me[1], date=date_today, days=5)
    for symbol in symbols:
      if not re.match(pattern, symbol):
        continue
      print(symbol)
      klr = klg._get_kline_request(date_today, f'{me[0]}:{me[1]}:{symbol}', '1m')
      kls = []
      for dst, ded in ranges:
        try:
          time.sleep(0.15)
          rets = query_history_kline_list([klr], dst, ded - datetime.timedelta(seconds = 1))
        except:
          break
        for ret in rets:
          kls.extend(klg.convert_klinelist(ret.klines)['klines'])
      if len(kls) == 0:
        continue
      df = pandas.DataFrame(kls)
      df['dti'] = pandas.DatetimeIndex(df['kline_timestamp'].astype(int))
      df.loc[df['dti'].dt.date == date_today.date()].reset_index(drop=True)
      df.index = df['dti']
      dfdict[symbol] = df
    df_multiindex = pandas.concat(dfdict, axis=1)
    df_multiindex.loc[:, df_multiindex.columns.get_level_values(1)=='open'].droplevel(1, axis=1)
    resolution_name = "PT1M"
    out_cache_dir = os.path.join(
        interval_klineh5_root,
        "coin",
        resolution_name,
        f"{market}.{exchange}",
        date_today.strftime("%Y%m%d"))
    os.makedirs(out_cache_dir, exist_ok=True)
    for col in ['open', 'close', 'high', 'low', 'volume']:
      dfcol = df_multiindex.loc[:, df_multiindex.columns.get_level_values(1)==col].droplevel(1, axis=1)
      dfcol.to_hdf(os.path.join(out_cache_dir, f'{col}.h5'), col, mode='w')


if __name__ == "__main__":
  fire.Fire(generate_klineh5)
